---
title: "Power Analysis"
author: "Manas GA"
date: "2024-11-14"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```

## Load packages

```{r message=FALSE, warning=FALSE}
library(standardize)
library(MCMCglmm)
library(ggplot2)
library(cowplot)
library(lmerTest)
library(latex2exp)
```

# Power analysis

The analysis output from simulated data is stored in null_distributions/PA_output.csv. This was generated by the script "3_Simulations.R".

We load this data, and create columns corresponding to sex and sex ratio. To do this, we first load our actual sex and sex ratio specific estimates of vA_w (from "Output/fitness_vA_output.csv"). We then fit a logistic regression to this data.

```{r}
# Load the actual estimates of vA_w (from "Output/fitness_vA_output.csv")
data_va = read.csv("Output/fitness_vA_output.csv", header=T)

test=FALSE

vA_w_female_fb = if(test){0.03}else{data_va[data_va$Sex=="Female Vw"&data_va$Sex.Ratio=="Female Biased",]$Va}
vA_w_female_mb = if(test){0.09}else{data_va[data_va$Sex=="Female Vw"&data_va$Sex.Ratio=="Male Biased",]$Va}
vA_w_male_fb = if(test){0.11}else{data_va[data_va$Sex=="Male Vw"&data_va$Sex.Ratio=="Female Biased",]$Va}
vA_w_male_mb = if(test){0.41}else{data_va[data_va$Sex=="Male Vw"&data_va$Sex.Ratio=="Male Biased",]$Va}

# Load output from analysed simulated data
d_sim = read.csv("null_distributions/PA_output.csv", header=T)

# Add columns for Sex and Sex.Ratio based on the values of vA_w in d_sim
d_sim$Sex = ifelse(abs(d_sim$vA_w-vA_w_female_fb)<1e-03|abs(d_sim$vA_w-vA_w_female_mb)<1e-03, "Female", "Male")
d_sim$Sex.Ratio = ifelse(abs(d_sim$vA_w-vA_w_female_fb)<1e-03|abs(d_sim$vA_w-vA_w_male_fb)<1e-03, "Female_biased", "Male_biased")

fit_power = glm(significance ~ COV_tw + nrep_t + vA_t + Sex + Sex.Ratio, d_sim, family = "binomial")

summary(fit_power)

d_sim$significance_pred = predict(fit_power, type = "response")

```

Next, to investigate how power changes as a function of sex, sex ratio, vA_t, and the number of replicate measurements for the trait, we plotted our fitted values using a function.

```{r}

plot_power_fitted = function(vA_t, Sex, Sex.Ratio){
  # Select data   
  d_sim_subset = d_sim[abs(d_sim$vA_t-vA_t)<1e-7&d_sim$Sex==Sex&d_sim$Sex.Ratio==Sex.Ratio,]
  
  if(nrow(d_sim_subset)==0){stop("This vA_t does not exist in the dataset or enter correct values of Sex ('Female' or 'Male') and Sex.Ratio ('Female_biased' or 'Female_biased')")}
  
  d_sim_subset$nrep_t = as.character(d_sim_subset$nrep_t)
  d_sim_subset$nrep_t = factor(d_sim_subset$nrep_t, levels=c("3", "10", "17", "24", "31"))


  plot_power = ggplot(d_sim_subset, aes(y = significance_pred*100, x =   COV_tw, color = nrep_t)) +
  theme_bw() + ylim(0, 100) +
  geom_point(size = 0.5) +
  labs(y = "Power (%)", x = "True Robertson Covariance", color = "Number of replicate measurements for the trait") +
  theme(text = element_text(size = 12.5)) + 
  ggtitle(paste(Sex, " fitness - ", sub("_", " ", Sex.Ratio), sep = ""), subtitle = TeX(paste("$V_{a,T} =$", vA_t, sep = ""))) + 
  theme(plot.title = element_text(hjust = 0.5, size = 12.5), plot.subtitle = element_text(hjust = 0.5, size = 12.5)) 
  
  return(plot_power)
  
}

##################
### Make plots ###
##################

plot_0.03_Female_fb = plot_power_fitted(0.03, "Female", "Female_biased") 
plot_0.03_Female_mb = plot_power_fitted(0.03, "Female", "Male_biased") 
plot_0.03_Male_fb = plot_power_fitted(0.03, "Male", "Female_biased") 
plot_0.03_Male_mb = plot_power_fitted(0.03, "Male", "Male_biased") 

plot_0.13_Female_fb = plot_power_fitted(0.13, "Female", "Female_biased") 
plot_0.13_Female_mb = plot_power_fitted(0.13, "Female", "Male_biased") 
plot_0.13_Male_fb = plot_power_fitted(0.13, "Male", "Female_biased") 
plot_0.13_Male_mb = plot_power_fitted(0.13, "Male", "Male_biased") 

plot_0.23_Female_fb = plot_power_fitted(0.23, "Female", "Female_biased")  
plot_0.23_Female_mb = plot_power_fitted(0.23, "Female", "Male_biased") 
plot_0.23_Male_fb = plot_power_fitted(0.23, "Male", "Female_biased") 
plot_0.23_Male_mb = plot_power_fitted(0.23, "Male", "Male_biased") 

plot_0.33_Female_fb = plot_power_fitted(0.33, "Female", "Female_biased") 
plot_0.33_Female_mb = plot_power_fitted(0.33, "Female", "Male_biased") 
plot_0.33_Male_fb = plot_power_fitted(0.33, "Male", "Female_biased") 
plot_0.33_Male_mb = plot_power_fitted(0.33, "Male", "Male_biased") 

plot_0.43_Female_fb = plot_power_fitted(0.43, "Female", "Female_biased")  
plot_0.43_Female_mb = plot_power_fitted(0.43, "Female", "Male_biased")  
plot_0.43_Male_fb = plot_power_fitted(0.43, "Male", "Female_biased") 
plot_0.43_Male_mb = plot_power_fitted(0.43, "Male", "Male_biased") 

# Capture the legend
legend_power = get_plot_component(plot_0.03_Female_fb + guides(color = guide_legend(nrow = 1, override.aes = list(size = 3))) + theme(legend.position = "bottom"),"guide-box-bottom")

# Make a combined plot without legends
plot_power_combined = plot_grid(plot_0.03_Female_fb + theme(legend.position = "none"),
                                plot_0.03_Female_mb + theme(legend.position = "none"),
                                plot_0.03_Male_fb + theme(legend.position = "none"),
                                plot_0.03_Male_mb + theme(legend.position = "none"),
                                plot_0.13_Female_fb + theme(legend.position = "none"),
                                plot_0.13_Female_mb + theme(legend.position = "none"),
                                plot_0.13_Male_fb + theme(legend.position = "none"),
                                plot_0.13_Male_mb + theme(legend.position = "none"),
                                plot_0.23_Female_fb + theme(legend.position = "none"),
                                plot_0.23_Female_mb + theme(legend.position = "none"),
                                plot_0.23_Male_fb + theme(legend.position = "none"),
                                plot_0.23_Male_mb + theme(legend.position = "none"),
                                plot_0.33_Female_fb + theme(legend.position = "none"),
                                plot_0.33_Female_mb + theme(legend.position = "none"),
                                plot_0.33_Male_fb + theme(legend.position = "none"),
                                plot_0.33_Male_mb + theme(legend.position = "none"),
                                plot_0.43_Female_fb + theme(legend.position = "none"),
                                plot_0.43_Female_mb + theme(legend.position = "none"),
                                plot_0.43_Male_fb + theme(legend.position = "none"),
                                plot_0.43_Male_mb + theme(legend.position = "none"),
                                labels = "AUTO",
                                ncol = 4)
plot_power_combined = plot_grid(plot_power_combined, legend_power, ncol = 1, rel_heights = c(1, 0.05))
ggsave(plot_power_combined,
       height = 15,
       width = 12,
       file = "Output/Figure_P1.jpg")
  

```

Next, we use the actual estimates of vA_w, vA_t, and the number of replicate measurements for each trait to predict statistical power as a function of the true RC. 
```{r}
# Load the trait dataframe ("Comprehensive_traits.csv") to list traits and get sample sizes

d = read.csv("Comprehensive_traits.csv", header = T)

d$Trait = factor(d$Trait)

# Make a list of all the traits  
list_traits = levels(d$Trait)

# Remove the names of fitness data.
list_traits = list_traits[grep("Fitness", list_traits, invert = T)]

# For traits measured at different sex ratios (eg, mate fecundity, mating rate) the sex-ratio suffix ("_Male_biased", etc) has to be removed for the application of prc

list_traits = gsub("_Female_biased", "", list_traits)
list_traits = gsub("_Equal", "", list_traits)
list_traits = gsub("_Male_biased", "", list_traits)

# This results in duplication of the names of traits measured at different sex ratios
# Selected only unique names

list_traits = unique(list_traits)

# Make a list of names of traits that should appear in plot titles

list_names = c("Copulation duration in females", "Copulation duration in males","Female fecundity", "Male effect on female fecundity", "Female latency to first mating", "Male latency to first mating", "Female mating rate", "Male mating rate", "Sperm defence ability (P1)", "Sperm offence ability (P2)", "Female remating latency", "Male mating latency with mated females")

print(cbind(list_names, list_traits))

# List of traits measured at sex ratio environments (also exists inside prc_function!!!!! If editing here, also edit there!!!)

sr_list = c("Female_fecundity","Mate_fecundity" , "MR_Male", "MR_Female")

# Load the data containing trait vA_w outputs

data_va_trait = read.csv("Output/trait_vA_output.csv", header=T)

# Load the data frames containing the actual RC estimates

data_rc = read.csv("Output/PRC_output.csv", header=T) # Non sex ratio traits
data_rc_sep_sr = read.csv("Output/PRC_output_sep_sr.csv", header=T) # Sex ratio traits

```
We then  write a function that plots predicted power as a function of true RC appropriately for sex ratio and non-sex ratio traits. This function has two arguments:
trait = a trait from list_traits
res = how many points to be used to make plots

If the trait is a sex ratio traits this function uses sex-ratio specific vA_t to calculate power to measure sex and sex-ratio specific RCs. Otherwise, the function uses just the one vA_t available for the trait.

```{r}
plot_power_trait = function(trait, res = 200){
  
  if(trait %in% sr_list){
    
    # Save the actual additive genetic variance for the trait
    
    vA_t_fb = data_va_trait[data_va_trait$trait==paste(trait, "_Female_biased", sep=""),]$va_trait_true
    vA_t_mb = data_va_trait[data_va_trait$trait==paste(trait, "_Male_biased", sep=""),]$va_trait_true
    
    # Save the actual Robertson Covariances (absolute)
    
    cov_trait_femaleW_fb = abs(as.numeric(sub(" .*", "", data_rc_sep_sr[data_rc_sep_sr$trait==trait,]$cov_trait_femaleW_fb)))
    cov_trait_femaleW_mb = abs(as.numeric(sub(" .*", "", data_rc_sep_sr[data_rc_sep_sr$trait==trait,]$cov_trait_femaleW_mb)))
    
    cov_trait_maleW_fb = abs(as.numeric(sub(" .*", "", data_rc_sep_sr[data_rc_sep_sr$trait==trait,]$cov_trait_maleW_fb)))
    cov_trait_maleW_mb = abs(as.numeric(sub(" .*", "", data_rc_sep_sr[data_rc_sep_sr$trait==trait,]$cov_trait_maleW_mb)))

    
    # Calculate the average umber of replicate measurements replicate measurements per hemigenome line for this trait
    d_trait = d[d$Trait==paste(trait, "_Female_biased", sep=""),]
    n_lines = length(unique(d_trait$Family))
    nrep_t = nrow(d_trait)/n_lines 
  }else{
    
    # Save the actual additive genetic variance for the trait
    
    vA_t_fb = data_va_trait[data_va_trait$trait==trait,]$va_trait_true
    vA_t_mb = data_va_trait[data_va_trait$trait==trait,]$va_trait_true
    
    # Save the actual Robertson Covariances
    
    cov_trait_femaleW_fb = abs(as.numeric(sub(" .*", "", data_rc[data_rc$trait==trait,]$cov_trait_femaleW_fb)))
    cov_trait_femaleW_mb = abs(as.numeric(sub(" .*", "", data_rc[data_rc$trait==trait,]$cov_trait_femaleW_mb)))
    
    cov_trait_maleW_fb = abs(as.numeric(sub(" .*", "", data_rc[data_rc$trait==trait,]$cov_trait_maleW_fb)))
    cov_trait_maleW_mb = abs(as.numeric(sub(" .*", "", data_rc[data_rc$trait==trait,]$cov_trait_maleW_mb)))

    
    # Calculate the average umber of replicate measurements replicate measurements per hemigenome line for this trait
    d_trait = d[d$Trait==trait,]
    n_lines = length(unique(d_trait$Family))
    nrep_t = nrow(d_trait)/n_lines
  }
  
  
  
  
  
  # Create new data frames (separate for each combination of sex and sex ratio) for which the power is to be predicted using fit_power
  
  ######################################
  ### Female fitness - female biased ###
  ######################################
  
  # sqrt(vA_t_fb*vA_w_female_fb)
  
  d_power_female_fb = data.frame("COV_tw" = seq(0, 0.35, length = res),
  "nrep_t" = nrep_t,
  "vA_t" = vA_t_fb,
  "Sex" = "Female",
  "Sex.Ratio" = "Female_biased")
  
  d_power_female_fb$significance_pred = predict(object = fit_power, newdata = d_power_female_fb, type = "response")
  
  plot_power_female_fb = ggplot(d_power_female_fb, aes(y = significance_pred*100, x =   COV_tw)) +
  theme_bw() + ylim(0, 100) +
  geom_point(size = 0.5) +
  labs(y = "Power (%)", x = "True Robertson Covariance") +
  theme(text = element_text(size = 12.5)) + 
  ggtitle(paste("Female fitness - Female biased")) + 
  theme(plot.title = element_text(hjust = 0.5, size = 12.5)) +
  geom_vline(xintercept = cov_trait_femaleW_fb, linetype = "dashed", color = "red")  
  
  # save the power to measure RC as large as the actual estimate
  
  power_femaleW_fb = 100*predict(object = fit_power,
                             newdata = data.frame("COV_tw" = cov_trait_femaleW_fb, "nrep_t" = nrep_t, "vA_t" = vA_t_fb, "Sex" = "Female", "Sex.Ratio" = "Female_biased"),
                             type = "response")
  
  ####################################
  ### Female fitness - male biased ###
  ####################################

  d_power_female_mb = data.frame("COV_tw" = seq(0, 0.35, length = res),
  "nrep_t" = nrep_t,
  "vA_t" = vA_t_mb,
  "Sex" = "Female",
  "Sex.Ratio" = "Male_biased")
  
  d_power_female_mb$significance_pred = predict(object = fit_power, newdata = d_power_female_mb, type = "response")
  
  plot_power_female_mb = ggplot(d_power_female_mb, aes(y = significance_pred*100, x =   COV_tw)) +
  theme_bw() + ylim(0, 100) +
  geom_point(size = 0.5) +
  labs(y = "Power (%)", x = "True Robertson Covariance") +
  theme(text = element_text(size = 12.5)) + 
  ggtitle(paste("Female fitness - Male biased")) + 
  theme(plot.title = element_text(hjust = 0.5, size = 12.5)) +
  geom_vline(xintercept = cov_trait_femaleW_mb, linetype = "dashed", color = "red")  
  
  # save the power to measure RC as large as the actual estimate
  
  power_femaleW_mb = 100*predict(object = fit_power,
                             newdata = data.frame("COV_tw" = cov_trait_femaleW_mb, "nrep_t" = nrep_t,"vA_t" = vA_t_mb, "Sex" = "Female", "Sex.Ratio" = "Male_biased"),
                             type = "response")
  
  
  ####################################
  ### Male fitness - female biased ###
  ####################################

  d_power_male_fb = data.frame("COV_tw" = seq(0, 0.35, length = res),
  "nrep_t" = nrep_t,
  "vA_t" = vA_t_fb,
  "Sex" = "Male",
  "Sex.Ratio" = "Female_biased")
  
  d_power_male_fb$significance_pred = predict(object = fit_power, newdata = d_power_male_fb, type = "response")
  
  plot_power_male_fb = ggplot(d_power_male_fb, aes(y = significance_pred*100, x =   COV_tw)) +
  theme_bw() + ylim(0, 100) +
  geom_point(size = 0.5) +
  labs(y = "Power (%)", x = "True Robertson Covariance") +
  theme(text = element_text(size = 12.5)) + 
  ggtitle(paste("Male fitness - Female biased")) + 
  theme(plot.title = element_text(hjust = 0.5, size = 12.5)) +
  geom_vline(xintercept = cov_trait_maleW_fb, linetype = "dashed", color = "red")  
  
  # save the power to measure RC as large as the actual estimate
  
  power_maleW_fb = 100*predict(object = fit_power,
                             newdata = data.frame("COV_tw" = cov_trait_maleW_fb, "nrep_t" = nrep_t, "vA_t" = vA_t_fb, "Sex" = "Male", "Sex.Ratio" = "Female_biased"),
                             type = "response")
  
  ##################################
  ### Male fitness - male biased ###
  ##################################

  d_power_male_mb = data.frame("COV_tw" = seq(0, 0.35, length = res),
  "nrep_t" = nrep_t,
  "vA_t" = vA_t_mb,
  "Sex" = "Male",
  "Sex.Ratio" = "Male_biased")
  
  d_power_male_mb$significance_pred = predict(object = fit_power, newdata = d_power_male_mb, type = "response")
  
  plot_power_male_mb = ggplot(d_power_male_mb, aes(y = significance_pred*100, x =   COV_tw)) +
  theme_bw() + ylim(0, 100) +
  geom_point(size = 0.5) +
  labs(y = "Power (%)", x = "True Robertson Covariance") +
  theme(text = element_text(size = 12.5)) + 
  ggtitle(paste("Male fitness - Male biased")) + 
  theme(plot.title = element_text(hjust = 0.5, size = 12.5)) +
  geom_vline(xintercept = cov_trait_maleW_mb, linetype = "dashed", color = "red")  
  
  # save the power to measure RC as large as the actual estimate
  
  power_maleW_mb = 100*predict(object = fit_power,
                             newdata = data.frame("COV_tw" = cov_trait_maleW_mb, "nrep_t" = nrep_t, "vA_t" = vA_t_mb, "Sex" = "Male", "Sex.Ratio" = "Male_biased"),
                             type = "response")
  
  ####################
  ### Common title ###
  ####################
  
  title = ggdraw() + draw_label(list_names[which(list_traits==trait)], hjust = 0.5, size = 18)
  
  # Combined plot without common title (trait)
  plot_power_trait_combined = plot_grid(plot_power_female_fb, 
                                        plot_power_female_mb,
                                        plot_power_male_fb,
                                        plot_power_male_mb,
                                        labels = "AUTO",
                                        nrow = 2)
  
  # Add title
  plot_power_trait_combined = plot_grid(title, plot_power_trait_combined, ncol = 1, rel_heights = c(0.1,1))
  actual_power = data.frame("trait" = trait, "power_femaleW_fb" = power_femaleW_fb, "power_femaleW_mb" = power_femaleW_mb, 
                            "power_maleW_fb" = power_maleW_fb, "power_maleW_mb" = power_maleW_mb)
  
  return(list(plot = plot_power_trait_combined, actual_power = actual_power))

  
}
```

Loop through the traits in list_traits and save the plots and actual power estimates generated by plot_power_trait()
```{r}

# Create an empty data frame

actual_power_data = data.frame()

for(trait in list_traits){
    
 # Get figure ID
 id = which(list_traits==trait) + 1
 
 power_analysis = plot_power_trait(trait = trait)
 figure = power_analysis$plot
 
 actual_power_data = rbind(actual_power_data, power_analysis$actual_power)
  
 ggsave(figure,
         width = 6,
         height = 6,
         file = paste("Output/Figure_P", id, ".jpg", sep=""))
  
}

actual_power_data$trait_names = list_names

write.table(actual_power_data, file = "Output/actual_power.csv",row.names = FALSE, sep = ",")
```

